[Ascend]Adapt to Most Operators #1002
Conversation
Summary of ChangesHello @qianjinqiu, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the Ascend backend's capabilities by integrating a broad array of new operators and optimizing existing ones through Triton kernel implementations. The changes aim to improve computational efficiency and expand the functional scope of the backend, ensuring better support for diverse deep learning models. Additionally, the update includes refinements to performance heuristics and standardization of debug logging for improved maintainability and clarity. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a significant number of new operators and fused kernels for the Ascend backend. While this is a great step towards feature parity, the implementation of several new Triton kernels contains critical correctness issues, such as using invalid syntax (Python loops inside JIT functions), which will prevent them from compiling. Additionally, there are multiple instances of severe performance anti-patterns, including launching kernels within loops and performing redundant multi-pass computations over the same data. I have also identified several maintainability concerns like dead code, non-English comments, and confusing logic. Addressing these issues is crucial for the stability, performance, and long-term health of the Ascend backend.
| if rotary_interleaved: | ||
| for d in range(0, BLOCK_D // 2): | ||
| dim_range_x = d * 2 | ||
| dim_range_y = d * 2 + 1 | ||
|
|
||
| rotary_embedding_rw_kernel( | ||
| state_out, | ||
| state, | ||
| cos, | ||
| sin, | ||
| stride_state_n, | ||
| stride_state_h, | ||
| stride_state_d, | ||
| stride_cos_n, | ||
| stride_cos_d, | ||
| num_tokens, | ||
| num_heads, | ||
| token_range, | ||
| head_range, | ||
| dim_range_x, | ||
| dim_range_y, | ||
| rotary_interleaved, | ||
| ) |
There was a problem hiding this comment.
The implementation for the rotary_interleaved=True case in rotary_embedding_siso_kernel is a major performance bottleneck. It uses a Python for loop to launch a separate kernel (rotary_embedding_rw_kernel) for each pair of dimensions. Launching kernels in a loop is a significant anti-pattern in Triton that leads to very poor performance due to high launch overhead. This logic should be vectorized to happen within a single kernel launch, similar to how the non-interleaved case is handled.
qianjinqiu
force-pushed
the
master
branch
8 times, most recently
from
6c8c939 to
b34f31e
Compare
qianjinqiu
force-pushed
the
master
branch
2 times, most recently
from
b05cd75 to
493dfad
Compare
2 participants
PR Category
Type of Change
Description
Issue
Progress
Performance